package skewreduce.framework.planner.partition;

import java.util.ArrayList;
import java.util.Arrays;
import java.util.Collections;
import java.util.Comparator;
import java.util.List;

import org.apache.hadoop.conf.Configurable;
import org.apache.hadoop.conf.Configuration;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import skewreduce.framework.planner.CostEstimator;
import skewreduce.framework.planner.EaggCostModel;
import skewreduce.framework.planner.PartitionNode;
import skewreduce.framework.planner.axis.Axis;
import skewreduce.framework.planner.partition.Partitioner.PartitionResult;
import skewreduce.framework.planner.partition.Partitioner.SplitResult;
import skewreduce.framework.planner.partition.Partitioner.TupleComparator;
import skewreduce.lib.ITuple;
import skewreduce.lib.Partition;
import skewreduce.lib.Util;


public abstract class PartitionStrategy implements Configurable {
	private static final Logger LOG = LoggerFactory.getLogger(PartitionStrategy.class);

	protected Configuration conf;
	protected EaggCostModel model;
	
	@Override
	public final Configuration getConf() {
		return conf;
	}
	
	@Override
	public void setConf(Configuration conf) {
		this.conf = conf;
	}
	
	public void initialize(EaggCostModel model) {
		this.model = model;
	}
	
	class PartitionContext {
		int pivot;
		double costL;
		double costR;
		double costMax;
		long mergeItems;
		
		PartitionContext() {
			reset();
		}
		
		public void set(int p,double l,double r,long m) {
			pivot = p; costL = l; costR = r; mergeItems = m;
			costMax = Math.max(l,r);
		}
		
		public int getPivot() {
			return pivot;
		}
		
		public boolean isSafe() {
			return mergeItems >= 0;
		}
		
		public double getLeftCost() { return costL; }
		public double getRightCost() { return costR; }
		public double getCost() { return costMax; }
		
		public void reset() {
			pivot = -1;
			costL = Double.MAX_VALUE;
			costR = Double.MAX_VALUE;
			costMax = Double.MAX_VALUE;
			mergeItems = Long.MIN_VALUE;
		}
	}
	
	protected PartitionResult splitPartition(PartitionNode parent,Axis axis,float split,boolean safe) {
		return splitPartition(parent,axis,split,0.0f,safe);
	}

	protected PartitionResult splitPartition(PartitionNode parent,Axis axis,float split, float prob,boolean safe) {
		int field = axis.getIndex();
		Partition[] children = parent.getPartition().split(field, split,prob);
        
        ArrayList<ITuple> ls = new ArrayList<ITuple>();
        ArrayList<ITuple> rs = new ArrayList<ITuple>();
        
        for ( ITuple t : parent.getSample() ) {
        	float v = t.getFloat(field);
        	if ( v < split ) {
        		ls.add(t);
        	} else if ( v > split ) {
        		rs.add(t);
        	} else if ( prob > 0.0f && Math.random() < prob ) {
        		ls.add(t);
        	} else {
        		rs.add(t);
        	}
        }
        
        ls.trimToSize();
        rs.trimToSize();
        
        /*
		if ( LOG.isDebugEnabled() ) {
			// do sanity check
			for ( ITuple t : ls ) {
				if ( t.getFloat(axis) > split ) {
					LOG.error("Wrong split!: {} - less than {} expected along axis {}",new Object[] { t, split, axis });
					throw new IllegalStateException();
				}
			}
			for ( ITuple t : rs ) {
				if ( t.getFloat(axis) < split ) {
					LOG.error("Wrong split!: {} - greater than {} expected along axis {}",new Object[] { t, split, axis });
					throw new IllegalStateException();
				}
			}
		} 
		*/       
        
        return new PartitionResult(safe,new PartitionNode(parent,children[0],ls), new PartitionNode(parent,children[1],rs));
	}
	
	protected int countLeft(List<ITuple> sample,int from,int field,float v) {
		int count = 0;
		for ( int i = from; i >= 0; --i ) {
			if ( sample.get(i).getFloat(field) == v ) {
				++count;
			} else {
				break;
			}
		}
		return count;
	}
	
	protected int countRight(List<ITuple> sample,int from,int field,float v) {
		int count = 0;
		for ( int i = from; i < sample.size(); ++i ) {
			if ( sample.get(i).getFloat(field) == v )
				++count;
			else
				break;
		}
		return count;
	}
	
	protected float computeSplitProb(List<ITuple> sample,int pivot,int field) {
		float split = sample.get(pivot).getFloat(field);
		int cntL = countLeft(sample,pivot-1,field,split);
		int cntR = countRight(sample,pivot+1,field,split);
		int total = cntL + cntR + 1; // including pivot
		return ( total == 1 ) ? 0.0f : cntL/(float)total;
	}
	
	protected PartitionResult splitSortedPartition(PartitionNode root,Axis axis,PartitionContext context) {
		return splitSortedPartition(root,axis,context.getPivot(),Float.NaN,Float.NaN,context.getLeftCost(),context.getRightCost(),context.isSafe());
	}
	
	protected PartitionResult splitSortedPartition(PartitionNode root,Axis axis,int pivot,double costL,double costR,boolean safe) {
		return splitSortedPartition(root,axis,pivot,Float.NaN,Float.NaN,costL,costR,safe);
	}
	
	protected PartitionResult splitSortedPartition(PartitionNode root,Axis axis,int pivot,float split,float prob,double costL,double costR,boolean safe) {
		List<ITuple> sample = root.getSample();
		int field = axis.getIndex();
		Partition parent = root.getPartition();
		
		List<ITuple> ls = sample.subList(0, pivot);
		List<ITuple> rs = sample.subList(pivot, sample.size());
		
		float bestProb = ( Float.isNaN(prob) ) ? computeSplitProb(sample,pivot,field) : prob;
		float bestSplit = ( Float.isNaN(prob) ) ? sample.get(pivot).getFloat(field) : split;
		
		Partition[] bestChildren = parent.split(field, bestSplit, bestProb);
		
		return new PartitionResult(safe,
				new PartitionNode(root,bestChildren[0],ls,costL,model.getWork2Sec()),
				new PartitionNode(root,bestChildren[1],rs,costR,model.getWork2Sec())
				);
	}
	
	protected PartitionResult splitSortedPartition2(TupleComparator comp,PartitionNode root,int pivot,double costL,double costR,boolean safe) {
		List<ITuple> sample = root.getSample();
		Partition parent = root.getPartition();
		
		List<ITuple> ls = sample.subList(0, pivot);
		List<ITuple> rs = sample.subList(pivot, sample.size());

		SplitResult res = comp.computeSplit(root.getSample(), pivot-1, pivot);
		
		Partition[] bestChildren = parent.split(res.getAxesOrder(), res.getSplit(), res.getLeftProb());
		
		return new PartitionResult(safe,
				new PartitionNode(root,bestChildren[0],ls,costL,model.getWork2Sec()),
				new PartitionNode(root,bestChildren[1],rs,costR,model.getWork2Sec())
				);
	}

	protected long isSafeForMerge(EaggCostModel model,Axis axis,List<ITuple> sample,float split) {
		if ( model.isMergeParallelizable() ) return 0; // don't care
		// check particles at merge surface does not exceed the limit
		int field = axis.getIndex();
		float r = (float)model.getBandWidth();
		float lowBound = split - r;
		float upBound = split + r;
		int nSample = 0;
		for ( ITuple t : sample ) {
			float v = t.getFloat(field);
			if ( lowBound <= v && v <= upBound ) ++nSample;
		}
		
		long expNumItems = (long)Math.ceil(nSample / model.getSampleRate());
		if ( LOG.isDebugEnabled() ) {
			LOG.debug("samples at merge = {}, expected items = {}, limit = {}",new Object[] { nSample,expNumItems,model.getMaxMergeDataItems() });
		}
		
		if ( expNumItems < model.getMaxMergeDataItems() )
			return expNumItems;
		else
			return -1 * expNumItems;
	}

	protected long isSafeForMergeForSorted(EaggCostModel model,Axis axis,List<ITuple> sample,int pivot,float split) {
		if ( model.isMergeParallelizable() ) return 0; // don't care
		// check particles at merge surface does not exceed the limit
		int field = axis.getIndex();
		float r = (float)model.getBandWidth();
		float lowBound = split - r;
		float upBound = split + r;

		int i = pivot;
		for ( ; i >= 0 && sample.get(i).getFloat(field) >= lowBound; --i );

		int j = pivot;
		for ( ; j < sample.size()-1 && sample.get(j).getFloat(field) <= upBound; ++j );
		
		int nSample = j - i + 1;
	
		long expNumItems = (long)Math.ceil((j - i + 1) / model.getSampleRate());
		if ( LOG.isDebugEnabled() ) {
			LOG.debug("samples at merge = {}, expected items = {}, limit = {}",new Object[] { nSample,expNumItems,model.getMaxMergeDataItems() });
		}
		
		if ( expNumItems < model.getMaxMergeDataItems() )
			return expNumItems;
		else
			return -1 * expNumItems;
	}
	
	public abstract PartitionResult partition(Axis axis,PartitionNode root,TupleComparator comp);
	

	// implementations
	
	public static final class BinarySearch extends PartitionStrategy {
		PartitionContext context = new PartitionContext();
		
		@Override
		public PartitionResult partition(Axis axis, PartitionNode root,TupleComparator comp) {
			List<ITuple> sample = root.getSample();
			Partition parent = root.getPartition();
			
			context.reset();
			
			int field = axis.getIndex();
			Collections.sort(sample,comp);

			int l = 0;
			int r = sample.size();
			do {
				int pivot = (l+r) >> 1;

				float split = sample.get(pivot).getFloat(field);
				
				// assign probability
				Partition[] children = parent.split(field, split);
				
				double costL = model.costWork(children[0], sample, 0, pivot); // L
				double costR = model.costWork(children[1], sample, pivot, sample.size()); // R
				double costMax = Math.max(costL, costR);

                if ( LOG.isDebugEnabled() ) {
                    LOG.debug("left={}/right={}",costL,costR);
                }
				
				if ( costMax < context.getCost() ) {
					long mitems = this.isSafeForMergeForSorted(model, axis, sample, pivot, split);
					
					boolean shouldKeep = ( ! context.isSafe() ) || mitems >= 0;
					if ( shouldKeep ) {
						context.set(pivot, costL, costR, mitems);
					}
				}
				
				if ( costL > costR ) {
					r = pivot-1;
				} else if ( costR > costL ) {
					l = pivot+1;
				} else {
					break;
				}
			} while ( l < r );
			
			return splitSortedPartition(root,axis,context);
		}
	}
	
	public static final class Sample extends PartitionStrategy {
		int numSample;
		PartitionContext context = new PartitionContext();
		
		@Override
		public void setConf(Configuration conf) {
			this.conf = conf;
			numSample = conf.getInt("skewreduce.planner.partition.sampler.numSample", 10);
		}
		
		@Override
		public PartitionResult partition(Axis axis, PartitionNode root, TupleComparator comp) {
			List<ITuple> sample = root.getSample();
			Partition parent = root.getPartition();
			
			context.reset();
			
			int field = axis.getIndex();
			Collections.sort(sample,comp);
			
			int stride = ( sample.size() > numSample ) ? sample.size() / numSample : 1;
			
			for ( int pivot = 1; pivot < sample.size(); pivot += stride ) {
				float split = sample.get(pivot).getFloat(field);
				Partition[] children = parent.split(field, split);

				double costL = model.costWork(children[0], sample, 0, pivot); // L
				double costR = model.costWork(children[1], sample, pivot, sample.size()); // R
				double costMax = Math.max(costL,costR);
				
				if ( costMax < context.getCost() ) {
					long mitems = this.isSafeForMergeForSorted(model, axis, sample, pivot, split);
					boolean shouldKeep = ( ! context.isSafe() ) || mitems >= 0;
					if ( shouldKeep ) {
						context.set(pivot, costL, costR, mitems);
					}
				}
			}
			
			return splitSortedPartition(root,axis,context);
		}
	}
	
	public static final class Incremental extends PartitionStrategy {
		private CostEstimator leftCost;
		private CostEstimator rightCost;
		
		@Override
		public void initialize(EaggCostModel model) {
			this.model = model;
			leftCost = model.createWorkCostEstimator();
			rightCost = model.createWorkCostEstimator();
		}

		@Override
		public PartitionResult partition(Axis axis, PartitionNode root, TupleComparator comp) {
			List<ITuple> samples = root.getSample();
			Partition parent = root.getPartition();
			
			int field = axis.getIndex();
			Collections.sort(samples,comp);

			int l = -1;
			int r = samples.size();
			// save cost up to index-1 samples
			double[] costVec = new double[samples.size()];

			if ( LOG.isDebugEnabled() ) {
				LOG.debug("# of samples = {}",samples.size());
			}
			
			leftCost.init(parent,samples.size() >> 1);
			rightCost.init(parent,samples.size() >> 1);
			
			leftCost.update(samples.get(++l));
			rightCost.update(samples.get(--r));
			
			float lv = samples.get(l).getFloat(field);
			float rv = samples.get(r).getFloat(field);
			int numLv = 1;
			int numRv = 1;
			
			double lCost = leftCost.getCost();
			double rCost = rightCost.getCost();
			
			costVec[l] = lCost;
			costVec[r] = rCost;
			
			// FIXME check boundary case. move left or right further and see if that improves the balance.
			
			while ( l+1 != r ) {
				boolean incrLeft;
				if ( lCost > rCost ) {
					incrLeft = false;
				} else if ( rCost > lCost ) {
					incrLeft = true;
				} else {
					incrLeft = Math.random() < 0.5;
				}
				
				if ( incrLeft ) {
					leftCost.update(samples.get(++l));
					lCost = costVec[l] = leftCost.getCost();
					float v = samples.get(l).getFloat(field);
					if ( v != lv ) {
						numLv = 0;
						lv = v;
					}
					++numLv;
				} else {
					rightCost.update(samples.get(--r));
					rCost = costVec[r] = rightCost.getCost();
					float v = samples.get(r).getFloat(field);
					if ( v != rv ) {
						numRv = 0;
						rv = v;
					}
					++numRv;
				}
			}
			
			if ( LOG.isDebugEnabled() ) {
				LOG.debug("initial: l={} (value={}, num={}, cost={}); r={} (value={}, num={}, cost={})",new Object[] {l,lv,numLv,lCost,r,rv,numRv,rCost});
			}
			
			// we have filled all samples. now cross the border and see whether it further reduces cost
			
			double prevBest = Math.max(lCost,rCost);
			float prob;
			float split;
			if ( lCost == rCost ) {
				// hah! the best case!
				if ( lv == rv ) {
					prob = (float)numLv / (numLv+numRv);
					split = lv;
				} else {
					split = ( lv + rv ) * 0.5f;
					prob = 0.0f;
				}
			} else {
				double lastCost = 0.0;
				if ( lCost > rCost ) {
					// try to advance right hand side
					int i;
					for ( i = l; i > 0; --i ) {
						rightCost.update(samples.get(i));
						lastCost = rightCost.getCost();
						double maxCost = Math.max(lastCost, costVec[i-1]);
						if ( prevBest < maxCost ) {
							break;
						} else {
							prevBest = maxCost;
						}
					}
					
					// how far did we advanced?
					int adjust = r - i - 1;
					if ( adjust > 0 ) {
						l -= adjust;
						r -= adjust;
						lCost = costVec[l];
						rCost = lastCost;
					}
				} else {
					// try to advance left hand side
					int i;
					for ( i = r; i < samples.size()-1; ++i ) {
						leftCost.update(samples.get(i));
						lastCost = leftCost.getCost();
						double maxCost = Math.max(lastCost, costVec[i+1]);
						if ( prevBest < maxCost ) {
							break;
						} else {
							prevBest = maxCost;
						}
					}
					
					int adjust = i - l - 1;
					if ( adjust > 0 ) {
						l += adjust;
						r += adjust;
						lCost = lastCost;
						rCost = costVec[r];
					}
				}
				
				lv = samples.get(l).getFloat(field);
				rv = samples.get(r).getFloat(field);
				if ( lv == rv ) {
					split = lv;
					prob = computeSplitProb(samples, r, field);
				} else {
					split = ( lv + rv ) * 0.5f;
					prob = 0.0f;
				}
			}
			
			if ( LOG.isDebugEnabled() ) {
				LOG.debug("final: l={} (value={}, cost={}); r={} (value={}, cost={}); split={}; prob={}",new Object[] {l, lv, lCost, r, rv, rCost, split, prob});
			}
			
			boolean isSafe = this.isSafeForMergeForSorted(model, axis, samples, l, lv) >= 0;
			
			leftCost.free();
			rightCost.free();
			
			//return this.splitSortedPartition(root, axis, r, split, prob, lCost, rCost, isSafe);
			return this.splitSortedPartition2(comp, root, r, lCost, rCost, isSafe);
		}
	}
	
	public static final class Median extends PartitionStrategy {
		@Override
		public PartitionResult partition(Axis axis, PartitionNode root, TupleComparator comp) {
			List<ITuple> samples = root.getSample();
			final int n = samples.size();
			int med = n >> 1;
			Partition parent = root.getPartition();
			
			Util.findKMedian(samples, med+1, 0, n-1, axis.getIndex());
			float split = samples.get(med).getFloat(axis.getIndex());
			float adjust = model.adjustSplit(axis, split);
			if ( adjust != 0 ) {
				if ( LOG.isDebugEnabled() )
					LOG.debug("adjust split value from {} to {}",split,split+adjust);
				split += adjust;
				// now repartition the samples according to new split value
				int rc;
				if ( adjust < 0 ) {
					rc = Util.partition(samples,0,med,split,axis.getIndex());
				} else {
					rc = Util.partition(samples, med+1, n-1, split, axis.getIndex());
				}
				if ( rc < 0 ) {
					LOG.error("Failed to partition?!");
				} else {
					med = rc;
				}
			}
			
			boolean isSafe = this.isSafeForMerge(model, axis, samples, split+adjust) >= 0;
			
			List<ITuple> ls = samples.subList(0, med);
			List<ITuple> rs = samples.subList(med, n);
			
			Partition[] children = parent.split(axis.getIndex(), split+adjust);

			return new PartitionResult(isSafe,new PartitionNode(root,children[0],ls), new PartitionNode(root,children[1],rs));
		}
	}
	
	public static final class Uniform extends PartitionStrategy {
		@Override
		public PartitionResult partition(Axis axis, PartitionNode root, TupleComparator comp) {
			float split = (float)axis.getMidPoint();
			boolean isSafe = this.isSafeForMerge(model, axis, root.getSample(), split) >= 0;
			return splitPartition(root,axis,split,isSafe); // distribute sample
		}
	}
}
